package com.hsoft.zhxy.common.filter;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Map;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;

import org.apache.commons.text.StringEscapeUtils;

import com.alibaba.fastjson.JSON;

import lombok.extern.log4j.Log4j2;

/**
 * All rights Reserved, Designed By 哈尔滨工业大学软件工程股份有限公司
 * @Title:        XssHttpServletRequestWrapper
 * @Description:  
 * @author: menglingwei
 * @date:   2020年5月17日 上午2:28:22
 */
@Log4j2
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    private static final String UTF_8 = "UTF-8";
    
    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    @Override
    public String getHeader(String name) {
        return StringEscapeUtils.escapeHtml4(super.getHeader(name));
    }

    @Override
    public String getQueryString() {
        return StringEscapeUtils.escapeHtml4(super.getQueryString());
    }

    @Override
    public String getParameter(String name) {
        return StringEscapeUtils.escapeHtml4(super.getParameter(name));
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        if(values != null) {
            int length = values.length;
            String[] escapseValues = new String[length];
            for(int i = 0; i < length; i++){
                escapseValues[i] = StringEscapeUtils.escapeHtml4(values[i]);
            }
            return escapseValues;
        }
        return values;
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        String str=getRequestBody(super.getInputStream());
        @SuppressWarnings("unchecked")
        Map<String,Object> map= JSON.parseObject(str,Map.class);
        Map<String,Object> resultMap=new HashMap<>(map.size());
        for(String key : map.keySet()){
            Object val=map.get(key);
            if(map.get(key) instanceof String){
                resultMap.put(key,StringEscapeUtils.escapeHtml4(val.toString()));
            }else{
                resultMap.put(key,val);
            }
        }
        str=JSON.toJSONString(resultMap);
        final ByteArrayInputStream bais = new ByteArrayInputStream(str.getBytes());
        return new ServletInputStream() {
            @Override
            public int read() throws IOException {
                return bais.read();
            }
            @Override
            public boolean isFinished() {
                return false;
            }
            @Override
            public boolean isReady() {
                return false;
            }
            @Override
            public void setReadListener(ReadListener listener) {
            }
        };
    }

    private String getRequestBody(InputStream stream) {
        String line = "";
        StringBuilder body = new StringBuilder();
        // 读取POST提交的数据内容
        BufferedReader reader = new BufferedReader(new InputStreamReader(stream, Charset.forName(UTF_8)));
        try {
            while ((line = reader.readLine()) != null) {
                body.append(line);
            }
        } catch (IOException e) {
            log.error("xss过滤器获取请求数据body时发生io错误", e);
        }
        return body.toString();
    }
}